from typing import List, Union

from ray.rllib.utils.framework import try_import_torch

torch, nn = try_import_torch()


def make_time_major(
    tensor: Union["torch.Tensor", List["torch.Tensor"]],
    *,
    trajectory_len: int = None,
    recurrent_seq_len: int = None,
):
    """Swaps batch and trajectory axis.

    Args:
        tensor: A tensor or list of tensors to swap the axis of.
            NOTE: Each tensor must have the shape [B * T] where B is the batch size and
            T is the trajectory length.
        trajectory_len: The length of each trajectory being transformed.
            If None then `recurrent_seq_len` must be set.
        recurrent_seq_len: Sequence lengths if recurrent.
            If None then `trajectory_len` must be set.

    Returns:
        res: A tensor with swapped axes or a list of tensors with
        swapped axes.
    """
    if isinstance(tensor, (list, tuple)):
        return [
            make_time_major(_tensor, trajectory_len, recurrent_seq_len)
            for _tensor in tensor
        ]

    assert (
        trajectory_len is not None or recurrent_seq_len is not None
    ), "Either trajectory_len or recurrent_seq_len must be set."

    # Figure out the sizes of the final B and T axes.
    if recurrent_seq_len is not None:
        assert len(tensor.shape) == 2
        # Swap B and T axes.
        tensor = torch.transpose(tensor, 1, 0)
        return tensor
    else:
        T = trajectory_len
        # Zero-pad, if necessary.
        tensor_0 = tensor.shape[0]
        B = tensor_0 // T
        if B != (tensor_0 / T):
            assert len(tensor.shape) == 1
            tensor = torch.cat(
                [
                    tensor,
                    torch.zeros(
                        trajectory_len - tensor_0 % T,
                        dtype=tensor.dtype,
                        device=tensor.device,
                    ),
                ]
            )
            B += 1

    # Reshape tensor (break up B axis into 2 axes: B and T).
    tensor = torch.reshape(tensor, [B, T] + list(tensor.shape[1:]))

    # Swap B and T axes.
    tensor = torch.transpose(tensor, 1, 0)

    return tensor


def vtrace_torch(
    *,
    target_action_log_probs: "torch.Tensor",
    behaviour_action_log_probs: "torch.Tensor",
    discounts: "torch.Tensor",
    rewards: "torch.Tensor",
    values: "torch.Tensor",
    bootstrap_values: "torch.Tensor",
    clip_rho_threshold: Union[float, "torch.Tensor"] = 1.0,
    clip_pg_rho_threshold: Union[float, "torch.Tensor"] = 1.0,
):
    r"""V-trace for softmax policies implemented with torch.

    Calculates V-trace actor critic targets for softmax polices as described in
    "IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner
    Architectures" by Espeholt, Soyer, Munos et al. (https://arxiv.org/abs/1802.01561)

    The V-trace implementation used here closely resembles the one found in the
    scalable-agent repository by Google DeepMind, available at
    https://github.com/deepmind/scalable_agent. This version has been optimized to
    minimize the number of floating-point operations required per V-Trace
    calculation, achieved through the use of dynamic programming techniques. It's
    important to note that the mathematical expressions used in this implementation
    may appear quite different from those presented in the IMPALA paper.

    The following terminology applies:
        - `target policy` refers to the policy we are interested in improving.
        - `behaviour policy` refers to the policy that generated the given
            rewards and actions.
        - `T` refers to the time dimension. This is usually either the length of the
            trajectory or the length of the sequence if recurrent.
        - `B` refers to the batch size.

    Args:
        target_action_log_probs: Action log probs from the target policy. A float32
            tensor of shape [T, B].
        behaviour_action_log_probs: Action log probs from the behaviour policy. A
            float32 tensor of shape [T, B].
        discounts: A float32 tensor of shape [T, B] with the discount encountered when
            following the behaviour policy. This will be 0 for terminal timesteps
            (done=True) and gamma (the discount factor) otherwise.
        rewards: A float32 tensor of shape [T, B] with the rewards generated by
            following the behaviour policy.
        values: A float32 tensor of shape [T, B] with the value function estimates
            wrt. the target policy.
        bootstrap_values: A float32 of shape [B] with the value function estimate at
            time T.
        clip_rho_threshold: A scalar float32 tensor with the clipping threshold for
            importance weights (rho) when calculating the baseline targets (vs).
            rho^bar in the paper.
        clip_pg_rho_threshold: A scalar float32 tensor with the clipping threshold
            on rho_s in \rho_s \delta log \pi(a|x) (r + \gamma v_{s+1} - V(x_s)).
    """
    log_rhos = target_action_log_probs - behaviour_action_log_probs

    rhos = torch.exp(log_rhos)
    if clip_rho_threshold is not None:
        clipped_rhos = torch.clamp(rhos, max=clip_rho_threshold)
    else:
        clipped_rhos = rhos

    cs = torch.clamp(rhos, max=1.0)
    # Append bootstrapped value to get [v1, ..., v_t+1]
    values_t_plus_1 = torch.cat(
        [values[1:], torch.unsqueeze(bootstrap_values, 0)], axis=0
    )

    deltas = clipped_rhos * (rewards + discounts * values_t_plus_1 - values)

    # Note: The original IMPALA code (and paper) suggested to perform the following
    # v-trace for-loop on the CPU, due to its sequential nature. However, modern GPUs
    # are quite optimized for these shorted for-loops, which is why it should be faster
    # nowadays to leave these operations on the GPU to avoid the GPU<>CPU transfer
    # penalty. This penalty can actually be quite massive on the LEarner actors, given
    # all other code is already well optimized.
    vs_minus_v_xs = [torch.zeros_like(bootstrap_values, device=deltas.device)]
    for i in reversed(range(len(discounts))):
        discount_t, c_t, delta_t = discounts[i], cs[i], deltas[i]
        vs_minus_v_xs.append(delta_t + discount_t * c_t * vs_minus_v_xs[-1])
    vs_minus_v_xs = torch.stack(vs_minus_v_xs[1:])

    # Reverse the results back to original order.
    vs_minus_v_xs = torch.flip(vs_minus_v_xs, dims=[0])

    # Add V(x_s) to get v_s.
    vs = torch.add(vs_minus_v_xs, values)

    # Advantage for policy gradient.
    vs_t_plus_1 = torch.cat([vs[1:], torch.unsqueeze(bootstrap_values, 0)], axis=0)
    if clip_pg_rho_threshold is not None:
        clipped_pg_rhos = torch.clamp(rhos, max=clip_pg_rho_threshold)
    else:
        clipped_pg_rhos = rhos
    pg_advantages = clipped_pg_rhos * (rewards + discounts * vs_t_plus_1 - values)

    # Make sure no gradients backpropagated through the returned values.
    return torch.detach(vs), torch.detach(pg_advantages)
